from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta
from api.sql_api import SQLExecutor
import json
from airflow.models import Variable

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

def execute_sql_task(**context):
    # 从Airflow Variables获取数据库连接信息
    db_config = {
        'host': Variable.get('HOLO_HOST'),
        'port': int(Variable.get('HOLO_PORT')),
        'database': Variable.get('HOLO_DATABASE'),
        'user': Variable.get('HOLO_USER'),
        'password': Variable.get('HOLO_PASSWORD')
    }
    
    # 获取需要执行的SQL列表
    sql_list = context['dag_run'].conf.get('sql_list', [])
    
    if not sql_list:
        raise ValueError("未提供SQL语句列表")
    
    # 初始化SQL执行器
    executor = SQLExecutor(**db_config)
    
    # 执行SQL并获取结果
    results = executor.execute_sql_batch(sql_list)
    
    # 将结果存储到XCom中
    context['task_instance'].xcom_push(key='sql_results', value=results)
    
    return results

with DAG(
    'sql_executor_dag',
    default_args=default_args,
    description='执行SQL语句批量任务',
    schedule_interval=None,
    catchup=False
) as dag:
    
    execute_task = PythonOperator(
        task_id='execute_sql_task',
        python_callable=execute_sql_task,
        provide_context=True,
    ) 